﻿/*
* Copyright (c) 2012 Kyle Hayward
* 
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to deal
* in the Software without restriction, including without limitation the rights
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
* copies of the Software, and to permit persons to whom the Software is
* furnished to do so, subject to the following conditions:
* 
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
* 
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
* THE SOFTWARE.
*/
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;

using OptixDotNet;
using OptixDotNet.Math;
using OptixDotNet.Utils;

using Optix = OptixDotNet;

namespace OptixDotNetUtils
{
    public delegate void ParseHandler( string line, string[] tokens );

    public struct ObjMaterial
    {
        public float Ns;
        public float Ni;
        public float d;
        public float Tr;
        public float illum;

        public Vector3 Tf;
        public Vector3 Ka;
        public Vector3 Kd;
        public Vector3 Ks;
        public Vector3 Ke;
    }

    public class ObjGroup
    {
        public string group;
        public string name;
        public string mtrl;
        public List<Int3> VIndices;
        public List<Int3> NIndices;
        public List<Int3> TIndices;

        public ObjGroup( string name, string group )
        {
            this.group = group;
            this.name = name;
            mtrl = null;
            VIndices = new List<Int3>();
            NIndices = new List<Int3>();
            TIndices = new List<Int3>();
        }

        public uint[] GetIndices()
        {
            uint[] indices = new uint[ VIndices.Count * 3 ];

            for ( int i = 0; i < VIndices.Count; i++ )
            {
                indices[ i * 3 + 0 ] = (uint)VIndices[ i ].X;
                indices[ i * 3 + 1 ] = (uint)VIndices[ i ].Y;
                indices[ i * 3 + 2 ] = (uint)VIndices[ i ].Z;
            }

            return indices;
        }
    }

    public class OBJLoader
    {
        public string           FileName { get; set; }

        public BoundingBox      BBox;

        public float            Scale = 1.0f;

        public bool             ParseMaterials = true;
        public bool             ParseNormals = true;
        public bool             GenerateNormals = true;
        public bool             GenerateGeometry = true;
        public bool             FrongFaceWinding = true;

        public List<Vector3>    Positions;
        public List<Vector3>    Normals;
        public List<Vector2>    Texcoords;

        protected Dictionary<string, ObjMaterial> mMtrls;
        protected string        mCurrentMtrlName;
        protected ObjMaterial   mCurrentMtrl;

        protected ObjGroup      mCurrentGroup = null;
        protected string        mCurrentGroupName;
        public List<ObjGroup>   Groups;

        protected string        mDirectory;

        public OBJLoader( string filename )
        {
            FileName = filename;

            if ( !File.Exists( filename ) )
            {
                throw new FileNotFoundException( "OBJLoader Error: File not found", filename );
            }

            Positions = new List<Vector3>();
            Normals = new List<Vector3>();
            Texcoords = new List<Vector2>();

            BBox = BoundingBox.Invalid;

            mMtrls = new Dictionary<string, ObjMaterial>();
            Groups = new List<ObjGroup>();
        }

        public void LoadContent()
        {
            Positions.Clear();
            Normals.Clear();
            Texcoords.Clear();

            Groups.Clear();
            mMtrls.Clear();

            mCurrentGroup = null;
            mCurrentMtrlName = null;
            mCurrentMtrl = default( ObjMaterial );

            Console.Write( "Loading: " + Path.GetFileName( FileName ) + "... " );
            float start = Time.GetTimeInSecs();

            mDirectory = Path.GetDirectoryName( FileName ) + "\\";

            mCurrentGroup = new ObjGroup( "default", "default" );
            Groups.Add( mCurrentGroup );

            ParseFile( ParseObjToken, FileName );
            CreateNormals();
            CreateGeometry();

            Console.WriteLine( "{0:0.00}s", Time.GetTimeInSecs() - start );
        }

        protected static char[] omitchars = new char[] { ' ', '\r', '\n', '\t', '/' };
        protected void ParseFile( ParseHandler parser, string filename )
        {
            StreamReader reader = new StreamReader( filename );

            while ( !reader.EndOfStream )
            {
                string line = reader.ReadLine();

                if ( string.IsNullOrEmpty( line ) )
                    continue;

                string[] tokens = line.Split( omitchars, StringSplitOptions.RemoveEmptyEntries );

                if ( tokens == null || tokens[ 0 ].StartsWith( "#" ) )
                    continue;

                parser( line, tokens );
            }
        }

        protected void ParseObjToken( string line, string[] tokens )
        {
            string token = tokens[ 0 ].ToLower();
            try
            {
                switch ( token )
                {
                    case "v":
                    {
                        float x = float.Parse( tokens[ 1 ] );
                        float y = float.Parse( tokens[ 2 ] );
                        float z = float.Parse( tokens[ 3 ] );

                        Vector3 p = new Vector3( x * Scale, y * Scale, z * Scale );
                        BBox.AddFloat3( ref p );

                        Positions.Add( p );
                        break;
                    }
                    case "vn":
                    {
                        if ( ParseNormals )
                        {
                            float x = float.Parse( tokens[ 1 ] );
                            float y = float.Parse( tokens[ 2 ] );
                            float z = float.Parse( tokens[ 3 ] );

                            Normals.Add( new Vector3( x, y, z ) );
                        }
                        break;
                    }
                    case "vt":
                    {
                        float x = float.Parse( tokens[ 1 ] );
                        float y = float.Parse( tokens[ 2 ] );

                        Texcoords.Add( new Vector2( x, y ) );
                        break;
                    }
                    case "f":
                    {
                        int v0 = 0, v1 = 0, v2 = 0;

                        if ( tokens.Length >= 10 )
                        {
                            //full v, vn, vt indices
                            v0 = int.Parse( tokens[ 1 ] );
                            v1 = int.Parse( tokens[ 4 ] );
                            v2 = int.Parse( tokens[ 7 ] );

                            mCurrentGroup.VIndices.Add( new Int3( v0 - 1, v1 - 1, v2 - 1 ) );

                            v0 = int.Parse( tokens[ 2 ] );
                            v1 = int.Parse( tokens[ 5 ] );
                            v2 = int.Parse( tokens[ 8 ] );

                            mCurrentGroup.TIndices.Add( new Int3( v0 - 1, v1 - 1, v2 - 1 ) );

                            if ( ParseNormals )
                            {
                                v0 = int.Parse( tokens[ 3 ] );
                                v1 = int.Parse( tokens[ 6 ] );
                                v2 = int.Parse( tokens[ 9 ] );

                                mCurrentGroup.NIndices.Add( new Int3( v0 - 1, v1 - 1, v2 - 1 ) );
                            }

                            //forget anything else
                        }
                        else if ( tokens.Length >= 7 )
                        {
                            //v, ( vn | vt ) indices
                            v0 = int.Parse( tokens[ 1 ] );
                            v1 = int.Parse( tokens[ 3 ] );
                            v2 = int.Parse( tokens[ 5 ] );

                            mCurrentGroup.VIndices.Add( new Int3( v0 - 1, v1 - 1, v2 - 1 ) );

                            v0 = int.Parse( tokens[ 2 ] );
                            v1 = int.Parse( tokens[ 4 ] );
                            v2 = int.Parse( tokens[ 6 ] );

                            if ( line.Contains( "//" ) )
                            {
                                if ( ParseNormals )
                                    mCurrentGroup.NIndices.Add( new Int3( v0 - 1, v1 - 1, v2 - 1 ) );
                            }
                            else
                            {
                                mCurrentGroup.TIndices.Add( new Int3( v0 - 1, v1 - 1, v2 - 1 ) );
                            }
                        }
                        else if ( tokens.Length >= 4 )
                        {
                            //v indices
                            v0 = int.Parse( tokens[ 1 ] );
                            v1 = int.Parse( tokens[ 2 ] );
                            v2 = int.Parse( tokens[ 3 ] );

                            mCurrentGroup.VIndices.Add( new Int3( v0 - 1, v1 - 1, v2 - 1 ) );
                        }

                        break;
                    }
                    case "g":
                    {
                        if ( !ParseMaterials )
                            break;

                        string name = tokens[ 1 ];
                        mCurrentGroupName = name;

                        mCurrentGroup = new ObjGroup( name, name );
                        Groups.Add( mCurrentGroup );
                        break;
                    }
                    case "usemtl":
                    {
                        if ( !ParseMaterials )
                            break;

                        string name = tokens[ 1 ];

                        string groupname = mCurrentGroupName + name;

                        mCurrentGroup = new ObjGroup( groupname, mCurrentGroupName );
                        Groups.Add( mCurrentGroup );

                        mCurrentGroup.mtrl = name;

                        break;
                    }
                    case "mtllib":
                    {
                        if ( !ParseMaterials )
                            break;

                        mCurrentMtrl = default( ObjMaterial );

                        string name = mDirectory + tokens[ 1 ];
                        ParseFile( ParseMtrlToken, name );

                        if ( mCurrentMtrlName != null )
                        {
                            mMtrls.Add( mCurrentMtrlName, mCurrentMtrl );
                        }
                        mCurrentMtrlName = null;
                        break;
                    }
                    default:
                        break;
                }
            }
            catch ( Exception ex )
            {
                Console.WriteLine( "Error parsing obj token: " + ex.Message );
            }
        }

        protected void ParseMtrlToken( string line, string[] tokens )
        {
            string token = tokens[ 0 ].ToLower();
            try
            {
                switch ( token )
                {
                    case "newmtl":
                    {
                        if ( mCurrentMtrlName != null )
                        {
                            mMtrls.Add( mCurrentMtrlName, mCurrentMtrl );
                        }

                        mCurrentMtrlName = tokens[ 1 ];
                        mCurrentMtrl = new ObjMaterial();
                        break;
                    }
                    case "ns":
                    {
                        mCurrentMtrl.Ns = float.Parse( tokens[ 1 ] );
                        break;
                    }
                    case "ni":
                    {
                        mCurrentMtrl.Ni = float.Parse( tokens[ 1 ] );
                        break;
                    }
                    case "d":
                    {
                        mCurrentMtrl.d = float.Parse( tokens[ 1 ] );
                        break;
                    }
                    case "tr":
                    {
                        mCurrentMtrl.Tr = float.Parse( tokens[ 1 ] );
                        break;
                    }
                    case "illum":
                    {
                        mCurrentMtrl.illum = float.Parse( tokens[ 1 ] );
                        break;
                    }
                    case "tf":
                    {
                        float x = float.Parse( tokens[ 1 ] );
                        float y = float.Parse( tokens[ 2 ] );
                        float z = float.Parse( tokens[ 3 ] );
                        mCurrentMtrl.Tf = new Vector3( x, y, z );
                        break;
                    }
                    case "ka":
                    {
                        float x = float.Parse( tokens[ 1 ] );
                        float y = float.Parse( tokens[ 2 ] );
                        float z = float.Parse( tokens[ 3 ] );
                        mCurrentMtrl.Ka = new Vector3( x, y, z );
                        break;
                    }
                    case "kd":
                    {
                        float x = float.Parse( tokens[ 1 ] );
                        float y = float.Parse( tokens[ 2 ] );
                        float z = float.Parse( tokens[ 3 ] );
                        mCurrentMtrl.Kd = new Vector3( x, y, z );
                        break;
                    }
                    case "ks":
                    {
                        float x = float.Parse( tokens[ 1 ] );
                        float y = float.Parse( tokens[ 2 ] );
                        float z = float.Parse( tokens[ 3 ] );
                        mCurrentMtrl.Ks = new Vector3( x, y, z );
                        break;
                    }
                    case "ke":
                    {
                        float x = float.Parse( tokens[ 1 ] );
                        float y = float.Parse( tokens[ 2 ] );
                        float z = float.Parse( tokens[ 3 ] );
                        mCurrentMtrl.Ke = new Vector3( x, y, z );
                        break;
                    }
                    case "map_ka":
                    {
                        break;
                    }
                    case "map_kd":
                    {
                        break;
                    }
                    case "map_ks":
                    {
                        break;
                    }
                    case "map_d":
                    {
                        break;
                    }
                    default:
                        break;
                }
            }
            catch ( Exception ex )
            {
                Console.WriteLine( "Error parsing mtrl token: " + ex.Message );
            }
        }

        protected virtual void CreateGeometry()
        {
            if ( GenerateGeometry == false )
                return;           
        }

        protected void CreateNormals()
        {
            //only allow normal generation for objs that don't have normals
            if ( Normals.Count > 0 || !GenerateNormals )
            {
                GenerateNormals = false;
                return;
            }

            float winding = 1.0f;
            if ( FrongFaceWinding == false )
                winding = -1.0f;

            Console.Write( "[Creating normals: " );
            float start = Time.GetTimeInSecs();

            //fill with zeros
            Vector3 zero = new Vector3( 0, 0, 0 );
            for ( int i = 0; i < Positions.Count; i++ )
                Normals.Add( zero );

            //loop through the triangles
            foreach ( ObjGroup group in Groups )
            {
                for ( int i = 0; i < group.VIndices.Count; i++ )
                {
                    Int3 index = group.VIndices[ i ];
                    Vector3 v0 = Positions[ index.X ];
                    Vector3 v1 = Positions[ index.Y ];
                    Vector3 v2 = Positions[ index.Z ];

                    Vector3 leg0 = v1 - v0;
                    Vector3 leg1 = v2 - v0;
                    Vector3 normal = winding * Vector3.Cross( ref leg0, ref leg1 );

                    Normals[ index.X ] += normal;
                    Normals[ index.Y ] += normal;
                    Normals[ index.Z ] += normal;
                }
            }

            for ( int i = 0; i < Positions.Count; i++ )
            {
                Normals[ i ].Normalize();
            }

            Console.Write( "{0:0.00}s] ", Time.GetTimeInSecs() - start );
        }

        protected bool ValidateGroup( ObjGroup group )
        {
            for ( int i = 0; i < group.VIndices.Count; i++ )
            {
                System.Diagnostics.Debug.Assert( group.VIndices[ i ].X < Positions.Count );
                System.Diagnostics.Debug.Assert( group.VIndices[ i ].Y < Positions.Count );
                System.Diagnostics.Debug.Assert( group.VIndices[ i ].Z < Positions.Count );

                if ( i < group.NIndices.Count && group.NIndices.Count > 0 && Normals.Count > 0 )
                {
                    System.Diagnostics.Debug.Assert( group.NIndices[ i ].X < Normals.Count );
                    System.Diagnostics.Debug.Assert( group.NIndices[ i ].Y < Normals.Count );
                    System.Diagnostics.Debug.Assert( group.NIndices[ i ].Z < Normals.Count );
                }

                if ( i < group.TIndices.Count && group.TIndices.Count > 0 && Texcoords.Count > 0 )
                {
                    System.Diagnostics.Debug.Assert( group.TIndices[ i ].X < Texcoords.Count );
                    System.Diagnostics.Debug.Assert( group.TIndices[ i ].Y < Texcoords.Count );
                    System.Diagnostics.Debug.Assert( group.TIndices[ i ].Z < Texcoords.Count );
                }
            }
            return true;
        }

        public void WriteExplodedOBJ( string path )
        {
            StreamWriter writer = new StreamWriter( path );

            int prevGroupVertCount = 0;
            int vertexCount = 0;
            foreach ( ObjGroup group in Groups )
            {
                //empty group
                if ( group.VIndices.Count == 0 && group.NIndices.Count == 0 && group.TIndices.Count == 0 )
                    continue;

                writer.WriteLine( "\n#\n# object {0}\n#", group.group );

                for ( int i = 0; i < group.VIndices.Count; i++ )
                {
                    Int3 index = group.VIndices[i];
                    writer.WriteLine( "v {0} {1} {2}", Positions[ index.X ].X, Positions[ index.X ].Y, Positions[ index.X ].Z );
                    writer.WriteLine( "v {0} {1} {2}", Positions[ index.Y ].X, Positions[ index.Y ].Y, Positions[ index.Y ].Z );
                    writer.WriteLine( "v {0} {1} {2}", Positions[ index.Z ].X, Positions[ index.Z ].Y, Positions[ index.Z ].Z );

                    vertexCount += 3;
                }
                writer.WriteLine( "# {0} vertices\n", group.VIndices.Count * 3 );

                for ( int i = 0; i < group.NIndices.Count; i++ )
                {
                    Int3 index = group.NIndices[ i ];
                    writer.WriteLine( "vn {0} {1} {2}", Normals[ index.X ].X, Normals[ index.X ].Y, Normals[ index.X ].Z );
                    writer.WriteLine( "vn {0} {1} {2}", Normals[ index.Y ].X, Normals[ index.Y ].Y, Normals[ index.Y ].Z );
                    writer.WriteLine( "vn {0} {1} {2}", Normals[ index.Z ].X, Normals[ index.Z ].Y, Normals[ index.Z ].Z );
                }
                writer.WriteLine( "# {0} normals\n", group.NIndices.Count * 3 );

                for ( int i = 0; i < group.TIndices.Count; i++ )
                {
                    Int3 index = group.TIndices[ i ];
                    writer.WriteLine( "vt {0} {1}", Texcoords[ index.X ].X, Texcoords[ index.X ].Y );
                    writer.WriteLine( "vt {0} {1}", Texcoords[ index.Y ].X, Texcoords[ index.Y ].Y );
                    writer.WriteLine( "vt {0} {1}", Texcoords[ index.Z ].X, Texcoords[ index.Z ].Y );
                }
                writer.WriteLine( "# {0} texture coordinates\n", group.TIndices.Count * 3 );

                //triangles
                writer.WriteLine( "g {0}", group.group );
                writer.WriteLine( "usemtl {0}", group.mtrl );

                for ( int i = prevGroupVertCount; i < vertexCount; i += 3 )
                {
                    if( group.TIndices.Count == 0 && group.NIndices.Count > 0 )
                    {
                        Int3 vIndex = group.VIndices[i];
                        Int3 nIndex = group.NIndices[i];

                        writer.WriteLine( "f {0}//{1} {2}//{3} {4}//{5}", vIndex.X + 1, nIndex.X + 1,
                                                                          vIndex.Y + 1, nIndex.Y + 1,
                                                                          vIndex.Z + 1, nIndex.Z + 1 );
                    }
                    else if ( group.NIndices.Count == 0 && group.TIndices.Count > 0 )
                    {
                        Int3 vIndex = group.VIndices[ i ];
                        Int3 tIndex = group.TIndices[ i ];

                        writer.WriteLine( "f {0}/{1} {2}/{3} {4}/{5}", vIndex.X + 1, tIndex.X + 1,
                                                                       vIndex.Y + 1, tIndex.Y + 1,
                                                                       vIndex.Z + 1, tIndex.Z + 1 );
                    }
                    else if ( group.NIndices.Count == 0 && group.TIndices.Count == 0 )
                    {
                        Int3 vIndex = group.VIndices[ i ];

                        writer.WriteLine( "f {0} {1} {2}", vIndex.X, vIndex.Y, vIndex.Z );
                    }
                    else
                    {
                        writer.WriteLine( "f {0}/{1}/{2} {3}/{4}/{5} {6}/{7}/{8}", i + 1, i + 1, i + 1,
                                                                                   i + 2, i + 2, i + 2,
                                                                                   i + 3, i + 3, i + 3 );
                    }
                }

                prevGroupVertCount = vertexCount;
            }

            writer.Close();
        }
    }
}
